import os, time
import pandas as pd
import numpy as np
from tqdm import tqdm
import cv2
import torch
import segmentation_models_pytorch as smp
from segmentation_models_pytorch.utils.functional import remove_dataparallel
from segmentation_models_pytorch.utils.functional import normalize

Wm = 640
Hm = 480
MASK_TH = 0.5
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)
pth_path = '/opt/data/private/project/adc_segmentation/ckpts/50UD_05300/erxian/Unet_se_resnext50_32x4d_latest_epoch.pth'
model = smp.Unet('se_resnext50_32x4d', encoder_weights=None, classes=1, activation=None)
# params = torch.load('pre_trained_model/model.pth', map_location=torch.device('cpu'))
params = torch.load(pth_path, map_location=torch.device('cpu'))
model.load_state_dict(params["state_dict"])
model.eval()
if torch.cuda.is_available():
    model.cuda()  # 把模型放进GPU加速，没有GPU请注释这一句

def run(img):
    [H, W, _] = img.shape
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 转换颜色通道，这一步非常重要，用cv2读取图片是BGR通道，模型需要的输入时RGB通道
    img = cv2.resize(img, (Wm, Hm), interpolation=cv2.INTER_LINEAR)  # 转换图片大小，这一步也非常重要
    img = normalize(img, mean, std, max_pixel_value=255.0)  # 归一化图片，这一步非常重要
    # 把img变成torch tensor 输入模型
    img = torch.from_numpy(np.moveaxis(img, -1, 0).astype(np.float32))
    img = torch.unsqueeze(img, 0)  # 输入模型img大小必须为 1 X 3 X 480 X 640

    # img 输入模型获得mask
    with torch.no_grad():
        if torch.cuda.is_available():
            output = model(img.cuda()).float()  # 有GPU时
        else:
            output = model(img).float()
    # output = model(img).float() # 没有GPU时
    result = torch.sigmoid(output)

    probability = result.detach().cpu().numpy()
    pred = np.squeeze(probability)
    mask = cv2.threshold(pred, MASK_TH, 1, cv2.THRESH_BINARY)[1]

    src_mask = cv2.resize(mask, (W, H), interpolation=cv2.INTER_NEAREST)
    return src_mask

def segImage(I,S,*args):
    debug = 0
    varargin = args
    boundaryWidth = 1
    boundaryColor = 'red'
    if debug:
        t0 = time.time()
    # Parse the optional inputs.
    if ( len(varargin) % 2 != 0 ):
            print('Extra Parameters passed to the function must be passed in pairs.')

    parameterCount = int(len(varargin)/2)

    for parameterIndex in range(parameterCount):
            parameterName = varargin[parameterIndex*2 - 1]
            parameterValue = varargin[parameterIndex*2]
            switcher = str.lower(parameterName)
            if switcher ==  'boundarycolor':
                boundaryColor = parameterValue
            elif switcher == 'boundarywidth':
                boundaryWidth = parameterValue
            else:
                print('Invalid parameter')
    
    # S = np.array(S, dtype = np.int16)
    S = np.int16(S)
    #S = im2db(S)
    [cy,cx] = np.gradient(S)
    if debug:
        t1 = time.time()
        print('gradient take time = {}'.format(t1 - t0))
    ccc = np.where((abs(cx)+abs(cy)) >0, 1, 0)
    ccc1 = np.where(ccc > 0, 0, 1)

    # ccc = ccc.astype(int)
    # ccc1 = ccc1.astype(int)
    ccc = np.uint8(ccc)
    ccc1 = np.uint8(ccc1)
    if boundaryWidth>1:
        boundaryWidth = np.ceil(boundaryWidth)
        dilateWindow = np.ones(boundaryWidth, boundaryWidth)
        ccc = cv2.dilate(ccc, dilateWindow, iterations=1) #imdilate(ccc,dilateWindow)
    elif boundaryWidth<1:
        print('boundaryWidth has been reset to 1.')

    if boundaryColor == 'red':
        I[:,:,0] = np.maximum(I[:,:,0],ccc)
        I[:,:,1] = np.minimum(I[:,:,1],ccc1)
        I[:,:,2] = np.minimum(I[:,:,2],ccc1)
    elif boundaryColor== 'black':
        I[:,:,0] = np.minimum(I[:,:,0],ccc1)
        I[:,:,1] = np.minimum(I[:,:,1],ccc1)
        I[:,:,2] = np.minimum(I[:,:,2],ccc1)
    else:
        print('Does not recognize boundaryColor other than red and black')
        
        I[:,:,0] = np.maximum(I[:,:,0],ccc1)
        I[:,:,1] = np.maximum(I[:,:,1],ccc1)
        I[:,:,2] = np.maximum(I[:,:,2],ccc1)
    if debug:
        t2 = time.time()
        print('segImg take time = {}'.format(t2 - t1))

    return I

path = r'/opt/data/private/project/adc_T9/50UD_05300/selected_train_data/Images'
dst_path = r'/opt/data/private/project/adc_T9/50UD_05300/selected_train_data/Images_add_mask'
for code in tqdm(os.listdir(path)):
    os.makedirs(os.path.join(dst_path, code), exist_ok=True)
    for img_name in tqdm(os.listdir(os.path.join(path, code))):
        img = cv2.imread(os.path.join(path, code, img_name), 1)
        mask = run(img)
        res = segImage(img/255, mask)
        cv2.imwrite(os.path.join(dst_path, code, img_name[:-3]+'png'), res * 255)
